﻿namespace CodedUIProviderWithMsTestContext.SpecFlowPlugin
{
    using System;
    using System.CodeDom;

    using MsTestContext.Generator.SpecflowPlugin;

    using TechTalk.SpecFlow.Generator;
    using TechTalk.SpecFlow.Generator.Plugins;
    using TechTalk.SpecFlow.Generator.UnitTestProvider;
    using TechTalk.SpecFlow.Utils;

    /// <summary>
    ///     The coded ui provider plugin.
    /// </summary>
    public class CodedUIProviderPlugin : IGeneratorPlugin
    {
        #region Public Methods and Operators

        /// <summary>
        /// The initialize.
        /// </summary>
        /// <param name="generatorPluginEvents">
        /// The generator plugin events.
        /// </param>
        /// <param name="generatorPluginParameters">
        /// The generator plugin parameters.
        /// </param>
        public void Initialize(
            GeneratorPluginEvents generatorPluginEvents,
            GeneratorPluginParameters generatorPluginParameters)
        {
            generatorPluginEvents.CustomizeDependencies += this.GeneratorPluginEventsOnCustomizeDependencies;
        }

        #endregion

        #region Methods

        private void GeneratorPluginEventsOnCustomizeDependencies(
            object sender,
            CustomizeDependenciesEventArgs customizeDependenciesEventArgs)
        {
            string unitTestProviderName =
                customizeDependenciesEventArgs.SpecFlowProjectConfiguration.GeneratorConfiguration
                    .GeneratorUnitTestProvider;

            if (unitTestProviderName.Equals("mstest", StringComparison.InvariantCultureIgnoreCase)
                || unitTestProviderName.Equals("mstest.2010", StringComparison.InvariantCultureIgnoreCase))
            {
                customizeDependenciesEventArgs.ObjectContainer
                    .RegisterTypeAs<CodedUIGeneratorProvider, IUnitTestGeneratorProvider>();
            }
        }

        #endregion
    }

    /// <summary>
    ///     The coded ui generator provider.
    /// </summary>
    public class CodedUIGeneratorProvider : MsTestContextGeneratorProvider
    {
        #region Constructors and Destructors

        /// <summary>
        /// Initializes a new instance of the <see cref="CodedUIGeneratorProvider"/> class.
        /// </summary>
        /// <param name="codeDomHelper">
        /// The code dom helper.
        /// </param>
        public CodedUIGeneratorProvider(CodeDomHelper codeDomHelper)
            : base(codeDomHelper)
        {
        }

        #endregion

        #region Public Methods and Operators

        /// <summary>
        /// The set test class.
        /// </summary>
        /// <param name="generationContext">
        /// The generation context.
        /// </param>
        /// <param name="featureTitle">
        /// The feature title.
        /// </param>
        /// <param name="featureDescription">
        /// The feature description.
        /// </param>
        public override void SetTestClass(
            TestClassGenerationContext generationContext,
            string featureTitle,
            string featureDescription)
        {
            base.SetTestClass(generationContext, featureTitle, featureDescription);

            foreach (CodeAttributeDeclaration declaration in generationContext.TestClass.CustomAttributes)
            {
                if (declaration.Name == "Microsoft.VisualStudio.TestTools.UnitTesting.TestClassAttribute")
                {
                    generationContext.TestClass.CustomAttributes.Remove(declaration);
                    break;
                }
            }

            generationContext.TestClass.CustomAttributes.Add(
                new CodeAttributeDeclaration(
                    new CodeTypeReference("Microsoft.VisualStudio.TestTools.UITesting.CodedUITestAttribute")));
        }

        #endregion
    }
}